package benchmark.energy.client;

import benchmark.client.DBProxy;
import benchmark.energy.EnergyDataTxGenerator;
import benchmark.energy.tx.*;
import benchmark.transaction.definition.AbstractTransaction;
import benchmark.utils.*;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import scala.Tuple6;

import java.io.IOException;
import java.sql.*;
import java.util.*;
import java.util.concurrent.*;

/**
 * Created by crusher. 2020.10.14
 */
public class MariaDBExecutorClient implements DBProxy {

    private ThreadPoolExecutor exe;
    private BlockingQueue<Connection> connectionPool = new LinkedBlockingQueue<>();
    private ListeningExecutorService service;
    private Map<Connection, Integer> connIdMap = new HashMap<>();
    private HashMap<String, Integer> name2Id = new HashMap<>();
    private HashMap<Integer, String> id2Name = new HashMap<>();
    // pair(id, (r_id, st_time, en_time, status, travel_t, seg_cnt))
    private ArrayList<Tuple6<Long, Long, Long, Integer, Integer, Integer>> buffer = new ArrayList<>();
    // pair(rid, (r_id, st_time, en_time, status, travel_t, seg_cnt))
    private HashMap<Long, Tuple6<Long, Long, Long, Integer, Integer, Integer>> lastRIdBuffer = new HashMap<>();
    private ExecutorService threadPool;

    public MariaDBExecutorClient(String serverHost, int parallelCnt, int queueLength) throws IOException, ClassNotFoundException, SQLException, InterruptedException {
        Class.forName("org.mariadb.jdbc.Driver");
        String initURL = "jdbc:mariadb://" + serverHost + ":3306";
        this.exe = new ThreadPoolExecutor(parallelCnt, parallelCnt, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(queueLength), (r, executor) -> {
            if (!executor.isShutdown()) {
                try {
                    executor.getQueue().put(r);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        });
        threadPool = Executors.newFixedThreadPool(10);
        exe.prestartAllCoreThreads();
        this.service = MoreExecutors.listeningDecorator(exe);
        for (int i = 0; i < parallelCnt; i++) {
            Connection conn = DriverManager.getConnection(initURL, "root", "root");
            this.connectionPool.offer(conn);
            connIdMap.put(conn, i);
        }
        buildMapBetweenIdAndName();
    }

    /**
     * create table test_temporal (n int,v float, st TIMESTAMP(0), en TIMESTAMP(0), PERIOD for time_period(st,en))
     * drop table test_temporal
     * insert test_temporal values(1, 3.4, '2021-06-23 21:56:18', '2021-06-24 20:03:17');
     * insert test_temporal values(1, 3.6, '2021-06-24 19:03:17', '2021-06-24 20:03:17');
     * select n, st, en, en-st from test_temporal where NOW() between st AND en and NOW()+600  between st AND en;
     * select n, v, st, en, UNIX_TIMESTAMP(en)-UNIX_TIMESTAMP(st), en-st from test_temporal;
     * -- and NOW() between st AND en
     * UPDATE test_temporal FOR PORTION OF time_period FROM '2021-06-24 20:13:17' TO NOW() SET v = 88 where n=1
     *
     * @param conn
     */
    private void initDB(Connection conn) {
        try (Statement stmt = conn.createStatement()) {
            conn.setAutoCommit(true);
            ResultSet rs = stmt.executeQuery("SELECT * FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = 're_europe'");
            if (rs.next()) return;
            stmt.execute("CREATE DATABASE re_europe");
            stmt.execute("USE re_europe");
            stmt.execute("CREATE table energy_node(raw_id int PRIMARY KEY, name varchar(255), country varchar(255), volt int, lat float, lon float");
            stmt.execute("CREATE table electric_line(r_start int, r_end int, x float, y float, r_limit int, r_length float)");
            stmt.execute("CREATE table temporal_status(node_id int, st_time TIMESTAMP(0), en_time TIMESTAMP(0), " +
                    "load_signal FLOAT, solar_cosmo_p FLOAT, solar_cosmo_u FLOAT, solar_ecmwf_p FLOAT, solar_ecmwf_u FLOAT, wind_cosmo_p FLOAT, wind_cosmo_u FLOAT, wind_ecmwf_p FLOAT, wind_ecmwf_u FLOAT, " +
                    "PERIOD FOR time_period(st_time, en_time))");
            // time would be [st_time, en_time]
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    @Override
    public String testServerClientCompatibility() throws UnsupportedOperationException {
        return null;
    }

    @Override
    public ListenableFuture<ServerResponse> execute(AbstractTransaction tx) {
        switch (tx.getTxType()) {
            case TX_IMPORT_STATIC_DATA:
                return this.service.submit(execute((ImportStaticDataTx) tx));
            case TX_ENTITY_TEMPORAL_CONDITION_QUERY:
                return this.service.submit(execute((EntityTemporalConditionTx) tx));
            case TX_QUERY_SNAPSHOT:
                return this.service.submit(execute((SnapshotQueryTx) tx));
            case TX_UPDATE_TEMPORAL_DATA:
                return this.service.submit(execute((UpdateTemporalDataTx) tx));
            case TX_QUERY_SNAPSHOT_AGGR_MAX:
                return this.service.submit(execute((SnapshotAggrMaxTx) tx));
            case TX_QUERY_SNAPSHOT_AGGR_DURATION:
                return this.service.submit(execute((SnapshotAggrDurationTx) tx));
            default:
                throw new UnsupportedOperationException();
        }
    }

    @Override
    public void createDB() throws IOException {
    }


    public void createIndexes() throws InterruptedException {
        Connection conn = connectionPool.take();

        try (PreparedStatement stmt0 = conn.prepareStatement("create index ind_rid_st_en on temporal_status(node_id, st_time, en_time)")) {
            conn.setAutoCommit(false);
            stmt0.execute();
            conn.commit();
        } catch (SQLException throwable) {
            throwable.printStackTrace();
        } finally {
            connectionPool.put(conn);
        }
    }

    @Override
    public void restartDB() throws IOException {
    }

    @Override
    public void shutdownDB() throws IOException {
    }

    private void buildMapBetweenIdAndName() throws InterruptedException {
        Connection conn = connectionPool.take();
        buildMapBetweenIdAndName(conn);
        connectionPool.put(conn);
        System.out.println("Build the map!");
    }

    private void buildMapBetweenIdAndName(Connection conn) {
        try (PreparedStatement stmt0 = conn.prepareStatement("select raw_id as node_id, name from energy_node")) {
            ResultSet rs = stmt0.executeQuery();
            while (rs.next()) {
                int id = rs.getInt("node_id");
                String name = rs.getString("name");
                id2Name.put(id, name);
                name2Id.put(name, id);
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void close() throws IOException, InterruptedException {
        try {
            service.shutdown();
            while (!service.isTerminated()) {
                service.awaitTermination(10, TimeUnit.SECONDS);
                long completeCnt = exe.getCompletedTaskCount();
                int remains = exe.getQueue().size();
                System.out.println(completeCnt + " / " + (completeCnt + remains) + " query completed.");
            }
            while (!connectionPool.isEmpty()) {
                Connection conn = connectionPool.take();
                conn.close();
            }
            System.out.println("Client exit. send " + exe.getCompletedTaskCount() + " lines.");
        } catch (SQLException e) {
            e.printStackTrace();
            throw new IOException(e);
        }
    }

    private Callable<ServerResponse> execute(ImportStaticDataTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                try (PreparedStatement stat1 = conn.prepareStatement("INSERT INTO energy_node VALUES (?,?,?,?,?)");
                     PreparedStatement stat2 = conn.prepareStatement("INSERT INTO electric_line VALUES (?,?,?,?,?,?)")) {
                    conn.setAutoCommit(false);
                    for (ImportStaticDataTx.EnergyNode p : tx.getNodes()) {
                        stat1.setInt(1, p.getRawId());
                        stat1.setString(2, p.getName());
                        stat1.setString(3, p.getCountry());
                        stat1.setInt(4, p.getVoltage());
                        stat1.setFloat(5, p.getLatitude());
                        stat1.setFloat(6, p.getLongitude());
                        stat1.addBatch();
                    }
                    stat1.executeBatch();
                    for (ImportStaticDataTx.EnergyRel r : tx.getRels()) {
                        stat2.setInt(1, r.getFromRawId());
                        stat2.setInt(2, r.getToRawId());
                        stat2.setFloat(3, r.getX());
                        stat2.setFloat(4, r.getY());
                        stat2.setInt(6, r.getLimit());
                        stat2.setFloat(7, r.getLength());
                        stat2.addBatch();
                    }
                    stat2.executeBatch();
                    conn.commit();
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                return new AbstractTransaction.Result();
            }
        };
    }

    public Callable<ServerResponse> execute(ImportTemporalDataTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                try (PreparedStatement stmt0 = conn.prepareStatement("insert into temporal_status(node_id,st_time,en_time,) values(?,?,?,?,?,?,?,?,?,?,?,?)")) {
                    conn.setAutoCommit(false);
                    for (ImportTemporalDataTx.EnergyStatus s : tx.data) {
                        stmt0.setInt(1, s.getNodeId());
                        stmt0.setTimestamp(2, new Timestamp(s.getTimePoint() * 1000L));
                        stmt0.setTimestamp(3, new Timestamp((s.getTimePoint() + 3600) * 1000L));
                        List<Float> vals = s.getStatus();
                        for (int i = 4; i < 4 + vals.size(); i++) {
                            stmt0.setFloat(i, vals.get(i - 4));
                        }
                        stmt0.addBatch();
                    }
                    stmt0.executeBatch();
                    conn.commit();
                } catch (SQLException e) {
                    e.printStackTrace();
                } finally {
                    conn.setAutoCommit(true);
                }
                return new AbstractTransaction.Result();
            }
        };
    }

    private Callable<ServerResponse> execute(UpdateTemporalDataTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                conn.setAutoCommit(true);
                StringBuilder sql = new StringBuilder();
                sql.append("UPDATE test_temporal FOR PORTION OF time_period FROM '")
                        .append(new Timestamp((long) tx.getStartTime() * 1000L)).append("' TO '")
                        .append(new Timestamp((long) tx.getEndTime() * 1000L)).append("' SET ");
                List<Float> vals = tx.getValues();
                for(int i = 0; i< EnergyDataTxGenerator.properties.length; i++) {
                    String propName = EnergyDataTxGenerator.properties[i];
                    Float val = vals.get(i);
                    sql.append(propName).append("=").append(val).append(',');
                }
                sql.setLength(sql.length()-1);
                sql.append(" WHERE node_id=").append(name2Id.get(tx.getEntityId()));
                try (PreparedStatement stmt0 = conn.prepareStatement(sql.toString())) {
                    stmt0.executeQuery();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
                return new AbstractTransaction.Result();
            }
        };
    }

    private Callable<ServerResponse> execute(SnapshotQueryTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                List<Pair<String, Float>> res = new ArrayList<>();
                String sql = "select node_id," + tx.getPropertyName() + " as property from temporal_status as ts " +
                        "where  \"" + new Timestamp((long) tx.getTimestamp() * 1000L) + "\" between st_time and en_time)";
                System.out.println(sql);
                try (PreparedStatement stat = conn.prepareStatement(sql)) {
                    ResultSet rs = stat.executeQuery();
                    if (rs.next()) {
                        res.add(Pair.of(id2Name.get(rs.getInt("node_id")), rs.getFloat("property")));
                    }
                }
                SnapshotQueryTx.Result result = new SnapshotQueryTx.Result();
                result.setStatus(res);
                return result;
            }
        };
    }


    private Callable<ServerResponse> execute(SnapshotAggrMaxTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                List<Pair<String, Float>> res = new ArrayList<>();
                String sql = "select node_id, max(" + tx.getP() + ") as max_p from temporal_status " +
                        "WHERE st_time<=\"" + new Timestamp((long) tx.getT1() * 1000L) + "\" and en_time >= \"" + new Timestamp((long) tx.getT0() * 1000L) + "\") group by node_id";
                System.out.println(sql);
                try (PreparedStatement stat = conn.prepareStatement(sql)) {
                    conn.setAutoCommit(false);
                    ResultSet rs = stat.executeQuery();
                    if (rs.next()) {
                        res.add(Pair.of(id2Name.get(rs.getInt("node_id")), rs.getFloat("max_p")));
                    }
                } catch (SQLException e) {
                    e.printStackTrace();
                }
                SnapshotAggrMaxTx.Result result = new SnapshotAggrMaxTx.Result();
                result.setPropMaxValue(res);
                return result;
            }
        };
    }

    private Callable<ServerResponse> execute(SnapshotAggrDurationTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                List<Triple<String, Integer, Integer>> res = new ArrayList<>();
                int t0 = tx.getT0(), t1 = tx.getT1();
                String sql = "select node_id, " + tx.getP() + " as property, st_time, en_time from temporal_status where st_time <= \"" +
                        new Timestamp((long) t1 * 1000L) + "\" and en_time >= \"" +
                        new Timestamp((long) t0 * 1000L) + "\") order by node_id, st_time";
                System.out.println(sql);
                TreeSet<Integer> intervalStarts = tx.getIntervalStarts();
                try (PreparedStatement stat = conn.prepareStatement(sql);
                     ResultSet rs = stat.executeQuery()) {
                    int lastId = -1;
                    Map<Integer, Integer> map = new HashMap<>();
                    while (rs.next()) {
                        int id = rs.getInt("node_id");
                        if(lastId!=id){
                            if(lastId!=-1){
                                for(Map.Entry<Integer, Integer> e : map.entrySet()){
                                    res.add(Triple.of(id2Name.get(lastId), e.getKey(), e.getValue()));
                                }
                            }
                            for(Integer intervalStart : intervalStarts){
                                map.put(intervalStart, 0);
                            }
                        }
                        float value = rs.getFloat("property");
                        int st = (int) (rs.getTimestamp("st_time").getTime() / 1000L);
                        int en = (int) (rs.getTimestamp("en_time").getTime() / 1000L);
                        int duration;
                        if (t0 >= st && t1 <= en) { // inner
                            duration = t1 - t0;
                        } else if (t0 >= st) { // left
                            duration = en - t0 + 1;
                        } else if (t1 >= en) { // middle
                            if (t1 == en) duration = t1 - st;
                            else duration = en - st + 1;
                        } else { // right
                            duration = t1 - st;
                        }
                        Integer valGrp = intervalStarts.floor((int)value);
                        if (valGrp != null) {
                            map.merge(valGrp, duration, Integer::sum);
                        }
                        lastId = id;
                    }
                }
                SnapshotAggrDurationTx.Result result = new SnapshotAggrDurationTx.Result();
                result.setStatusDuration(res);
                return result;
            }
        };
    }

    private Callable<ServerResponse> execute(EntityTemporalConditionTx tx) {
        return new Req() {
            @Override
            protected AbstractTransaction.Result executeQuery(Connection conn) throws Exception {
                List<String> res = new ArrayList<>();
                String sql = "select node_id from temporal_status where st_time <= \"" + new Timestamp((long) tx.getT1() * 1000L) +
                        "\" and en_time >= \"" + new Timestamp((long) tx.getT0() * 1000L) + "\" and " + tx.getP() +
                        " between " + tx.getVMin() + " and " + tx.getVMax() + " limit 1";
                System.out.println(sql);
                try (PreparedStatement stat = conn.prepareStatement(sql)) {
                    ResultSet rs = stat.executeQuery();
                    if (rs.next()) {
                        res.add(String.valueOf(rs.getInt("node_id")));
                    }
                }
                EntityTemporalConditionTx.Result result = new EntityTemporalConditionTx.Result();
                result.setEntities(res);
                return result;
            }
        };
    }

    private abstract class Req implements Callable<ServerResponse> {
        private final TimeMonitor timeMonitor = new TimeMonitor();
        final AbstractTransaction.Metrics metrics = new AbstractTransaction.Metrics();

        private Req() {
            timeMonitor.begin("Wait in queue");
        }

        @Override
        public ServerResponse call() throws Exception {
            try {
                Connection conn = connectionPool.take();
                timeMonitor.mark("Wait in queue", "query");
                AbstractTransaction.Result result = executeQuery(conn);
                timeMonitor.end("query");
                if (result == null) throw new RuntimeException("[Got null. Server close connection]");
                connectionPool.put(conn);
                metrics.setWaitTime(Math.toIntExact(timeMonitor.duration("Wait in queue")));
                metrics.setSendTime(timeMonitor.beginT("query"));
                metrics.setExeTime(Math.toIntExact(timeMonitor.duration("query")));
                metrics.setConnId(connIdMap.get(conn));
                ServerResponse response = new ServerResponse();
                response.setMetrics(metrics);
                response.setResult(result);
                return response;
            } catch (Exception e) {
                e.printStackTrace();
                throw e;
            }
        }

        protected abstract AbstractTransaction.Result executeQuery(Connection conn) throws Exception;
    }
}

